module Structure

using Printf
using Utils

export GeneralAtom, to_str_pdb, to_str_gro

abstract type AbstractCoordinateItem
end

struct GeneralAtom <: AbstractCoordinateItem
    atname::String              # atom name
    atid::Int32                 # atom number
    resname::String             # residue name
    resid::Int32                # residue number
    chainname::String           # chain name
    chainid::Int32              # chain number
    x::Float64                  # coordinate
    y::Float64                  # coordinate
    z::Float64                  # coordinate
    vx::Float64                 # velocity
    vy::Float64                 # velocity
    vz::Float64                 # velocity
    typ::String                 # PDB only, ATOM or HETATM
    altloc::String              # PDB only
    icode::String               # PDB only
    occupy::Float64             # PDB only
    bf::Float64                 # PDB only, can write comment infomation
    element::String             # element name
    charge_formal::String       # atom formal charge
    attype::String              # pdbqt only, atom type
    charge_partial::Float64     # pdbqt only, atom partial charge
end

function GeneralAtom(; atname="", atid=0, resname="", resid=0, chainname="", chainid=0, x=0.0, y=0.0, z=0.0, vx=0.0, vy=0.0, vz=0.0,
    typ="", altloc="", icode="", occupy=1.0, bf=0.0, element="", charge_formal="", attype="", charge_partial=0.0)
    GeneralAtom(atname, atid, resname, resid, chainname, chainid, x, y, z, vx, vy, vz, typ, altloc, icode, 
    occupy, bf, element, charge_formal, attype, charge_partial)
end

function to_str_pdb(g_atom::GeneralAtom)
    if g_atom.typ == "ATOM" || g_atom.typ == "HETATM"
        s = @sprintf("%-6s%5d %-4s%1s%3s %1s%4d%1s   %8.3f%8.3f%8.3f%6.2f%6.2f          %2s",
                g_atom.typ, g_atom.atid, atid_for_pdb(g_atom.atname), g_atom.altloc, g_atom.resname, g_atom.chainname, g_atom.resid, 
                g_atom.icode, g_atom.x * 10, g_atom.y * 10, g_atom.z * 10, g_atom.occupy, g_atom.bf, g_atom.element)
        if !isempty(g_atom.charge_formal)
            s *= @sprintf("%-2s\n", g_atom.charge_formal)
        else
            s *= @sprintf("  \n")
        end
        return s
    elseif g_atom.typ == "TER"
        s = @sprintf("%-6s", g_atom.typ)
        if g_atom.atid != 0
            s *= @sprintf("%5d      %3s %1s%4d%1s\n", g_atom.atid, g_atom.resname, g_atom.chainname, g_atom.resid, g_atom.icode)
        else
            s *= @sprintf("                     \n")
        end
        return s
    end
end

function to_str_gro(g_atom::GeneralAtom)
    s = @sprintf("%5d%-3s%7s%5d%8.3f%8.3f%8.3f", g_atom.resid, g_atom.resname, g_atom.atname, g_atom.atid, g_atom.x, g_atom.y, g_atom.z)
    if g_atom.vx != 0 && g_atom.vy != 0 && g_atom.vz != 0
        s *= @sprintf("%8.3f%8.3f%8.3f\n", g_atom.vx, g_atom.vy, g_atom.vz)
    else
        s *= "\n"
    end
    return s
end

function atid_for_pdb(atid::String)
    # 从方便理解的atid到pdb使用的atid
    if length(atid) == 4
        if match(r"[A-Z][ABGDEZH0-9][1-9'][1-9]", atid) !== nothing
            return atid[4] * atid[1:3]
        else
            warning_hint("Warning: atom name may be unstandard: $atid\n")
            return atid
        end
    else
        return " " * atid
    end
end

end # module
